// 使用opencv模块部署yolov5-6.0版本

#include "yolo.h"
#include <iostream>
#include<opencv2/opencv.hpp>
#include<math.h>

using namespace std;
using namespace cv;
using namespace dnn;

int main()
{	
	cout<<"start-------------------"<<endl;
	string img_path = "/home/lin/Pictures/5.jpg";
	string model_path = "/home/lin/code/object detection/yolov5/runs/train/exp3/weights/best.onnx";
	//int num_devices = cv::cuda::getCudaEnabledDeviceCount();
	//if (num_devices <= 0) {
		//cerr << "There is no cuda." << endl;
		//return -1;
	//}
	//else {
		//cout << num_devices << endl;
	//}

	Yolo test;
	Net net;
	cout<<test.readModel(net, model_path, false)<<endl;
	if (test.readModel(net, model_path, false)) {
		cout << "read net ok!" << endl;
	}
	else {
		cout << "read net fail!" << endl;
		return -1;

	}

	//生成随机颜色
	vector<Scalar> color;
	srand(time(0));
	for (int i = 0; i < 80; i++) {
		int b = rand() % 256;
		int g = rand() % 256;
		int r = rand() % 256;
		color.push_back(Scalar(b, g, r));
	}
// ------------图片检测--------------------
	// vector<Output> result;
	// Mat img = imread(img_path);
	
	// if (test.Detect(img, net, result)) {
	// 	test.drawPred(img, result, color);

	// }
	// else {
	// 	cout << "Detect Failed!"<<endl;
	// }

	// waitKey(0);

// ---------------视频检测----------------
//	VideoCapture cap(0);
	VideoCapture cap("/home/lin/code/object detection/img2video/build/test.avi");
	Mat img;

	namedWindow("detect output",WINDOW_NORMAL);
     // resizeWindow("detect output", 420, 360);
      moveWindow("detect output", 1300, 0);
 

	while(cap.isOpened()){
		cap >> img;
		vector<Output> result;
		vector<Output> result_sort_x;
		imshow("src",img);
		
		if (test.Detect(img, net, result)) {
			// test.drawPred(img, result, color);
			
			result_sort_x = test.Sort(result);
			test.drawPred(img, result_sort_x, color);

		}
		imshow("detect output",img);
		// else {
		// 	cout << "Detect Failed!"<<endl;
		// }
		if(waitKey(20)==27){
			break;
		}
	}
	cap.release();
	destroyAllWindows();
    return 0;
}

